package bigo;

import bigo.data.ArrayVector;
import bigo.data.Matrix;
import bigo.data.MatrixMeta;
import bigo.data.Vector;
import bigo.data.Vector.Wrapper;
import bigo.lib.BigOReducer;
import bigo.lib.Combiner;
import bigo.lib.InputFormat;
import bigo.lib.OutputFormat;
import bigo.lib.Reducer;
import bigo.lib.TextMatrixInputFormat;
import bigo.lib.TextMatrixOutputFormat;
import java.io.IOException;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FsShell;
import org.apache.hadoop.io.IntWritable;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.util.ToolRunner;

/**
 * This Driver is used for implement the very basic KMeans implementation
 * @author Song Liu (sl9885@bristol.ac.uk)
 */
public class KMeansDriver extends LargeASmallBMultiplicationDriver {

    @InputFormat
    public static Class input = TextMatrixInputFormat.class;
    @OutputFormat
    public static Class output = TextMatrixOutputFormat.class;

    public KMeansDriver() {
    }

    /**
     * Construct A KMeansDriver with A(Data Vector), B(the initial centroids),
     * C the intermediate output
     * @param A Data Vector
     * @param B initial centroids
     * @param C intermediate output
     */
    public KMeansDriver(MatrixMeta A, MatrixMeta B, MatrixMeta C) {
        super(A, B, C);
    }

    /**
     * This Mapper computes distance between each data vector and possible centroids,
     * and output the Vector and the index of the nearest centroid.
     */
    @bigo.lib.Mapper
    public static class KMeansMapper extends LargeASmallBMultiplicationDriver.Mapper {

        @Override
        protected void map(IntWritable key,
                Vector.Wrapper value, Context context)
                throws IOException, InterruptedException {
            Vector v = (Vector) value.get();
            ArrayVector distances = computeDistances(v, multiplier);
            //find the smallest distance
            float smallest = Float.MAX_VALUE;
            int smallestIdx = -1;
            for (int i = 0; i < distances.length(); i++) {
                if (distances.get(i) < smallest) {
                    smallest = distances.get(i);
                    smallestIdx = i;
                }
            }

            //output
            context.write(new IntWritable(smallestIdx), value);
        }

        /**
         * Get the Euclidean distance
         * @param v
         * @param m
         * @return
         */
        public ArrayVector computeDistances(Vector v, Matrix m) {
            ArrayVector ret = new ArrayVector(m.getRowVectorAt(m.getColumnOffset()).length());

            int startFrom = m.getColumnOffset();
            for (int i = startFrom; i < startFrom + m.getRowArrays().size(); i++) {
                float sum = 0;
                for (Integer j : m.getRowVectorAt(i).keySet()) {
                    if (v.get(j) != null) {
                        sum += Math.pow(v.get(j) - m.getRowVectorAt(i).get(j), 2);
                    }
                }
                ret.put(i, (float) Math.sqrt(sum));
            }
            return ret;
        }
    }

    /**
     * Average the vectors that near each centroid and get the new centroid
     */
    @Combiner
    @Reducer
    public static class AggregationReducer extends BigOReducer {

        @Override
        protected void reduce(IntWritable key, Iterable<Wrapper> values,
                Context context) throws IOException, InterruptedException {
            context.write(key, new Wrapper(centroid(values)));
        }

        /**
         * Simple centroid algorithm
         * @param values
         * @return
         */
        public Vector centroid(Iterable<Wrapper> values) {
            int num = 0;

            ArrayVector ret = null;
            for (Wrapper wrapper : values) {
                if (ret == null) {
                    ret = new ArrayVector(((Vector) wrapper.get()).length());
                }
                for (Integer i : ((Vector) wrapper.get()).keySet()) {
                    ret.put(i, ret.get(i) + ((Vector) wrapper.get()).get(i));
                }
                num++;
            }

            for (int i = 0; i < ret.length(); i++) {
                ret.put(i, ret.get(i) / num);
            }
            return ret;
        }
    }

    @Override
    public void init(Job job) {
        super.init(job);
    }

    @Override
    public int run(String[] args) throws Exception {
        int m = Integer.parseInt(args[2]), k = Integer.parseInt(args[3]),
                n = Integer.parseInt(args[4]);
        KMeansDriver.A = new MatrixMeta(args[0], m, k);
        KMeansDriver.B = new MatrixMeta(args[1], k, n);
        KMeansDriver.C = new MatrixMeta("output", m, n);

        int res = super.run(args);
        if (res == 0) {
            //download result to local
            FsShell fs = new FsShell(getConf());
            new ProcessBuilder("rm", "-r", args[1]).start().waitFor();
            fs.run(new String[]{"-get", "output", args[1]});
        }
        return res;
    }

    public static void main(String[] args) throws Exception {
        ToolRunner.run(new Configuration(), new KMeansDriver(), args);
    }
}
